import java.io.*;
import java.net.*;
import java.nio.*;
import java.util.*;

class TCPClient extends TCP
{
	public static void main(String argv[]) throws Exception
	{
		byte[] encryptionKey = null;
		byte[] integrityKey = null;
		byte[] IV = null;
		
		String opt = "";
		while (true)
		{
			System.out.println();
			System.out.println("1. Start session with Bob");
			System.out.println("2. Resume session with Bob");
			System.out.println("3. Send file to Bob");
			System.out.println("4. Request file from Bob");
			System.out.println("5. Exit\n");
			opt = Utils.GetUserInput("Enter selection: ");
			
			if (opt.equals("1"))
			{
				if (socket != null && !socket.isClosed())
				{
					System.out.println("There is already an active session");	
					continue;
				}
				
				socket = new Socket("localhost", 5555);
				out = new DataOutputStream(socket.getOutputStream());
				in = new DataInputStream(socket.getInputStream());
				
				if (startSession())
				{
					System.out.println("Session established successfully");
					encryptionKey = Utils.HMAC(k, new byte[] {0x01});
					integrityKey = Utils.HMAC(k, new byte[] {0x02});
					IV = Arrays.copyOfRange(Utils.HMAC(k, new byte[]{0x03}), 0, 8);
				}
				else
				{
					sessionId = 0;
					in.close();
					out.close();
					socket.close();
					System.out.println("Session initiation failed");
				}
			}
			
			else if (opt.equals("2"))
			{
				if (sessionId == 0)
					System.out.println("There is no current active session");
				else if (!socket.isClosed())
					System.out.println("There is already an active session");	
				else
				{
					socket = new Socket("localhost", 5555);
					out = new DataOutputStream(socket.getOutputStream());
					in = new DataInputStream(socket.getInputStream());
					
					if (resumeSession())
					{
						System.out.println("Session established successfully");
						encryptionKey = Utils.HMAC(k, new byte[] {0x01});
						integrityKey = Utils.HMAC(k, new byte[] {0x02});
						IV = Arrays.copyOfRange(Utils.HMAC(k, new byte[]{0x03}), 0, 8);
					}
					else
					{
						in.close();
						out.close();
						socket.close();
						System.out.println("Session resume failed");
					}
				}
			}
			
			else if (opt.equals("3"))
			{
				if (sessionId == 0 || socket.isClosed())
					System.out.println("There is no current active session");
				else
				{				
					String fileSendName = Utils.GetUserInput("Enter filename to send: "); // or can be filename requested from Bob
					byte[] fileBytes = Utils.getFileBytes(fileSendName);
					EncryptClass sendFile = new EncryptClass(cipher.toString(), encryptionKey, integrityKey, IV, fileBytes, true);
					byte[] sendBytes = sendFile.getCipher();
					
					Message m = new Message(Message.MessageType.FileSend);
					m.data = sendBytes;
					
					SendMessage(SimulateAttacker(m));
					
					in.close();
					out.close();
					socket.close();
				}
			}
			
			else if (opt.equals("4"))
			{
				if (sessionId == 0 || socket.isClosed())
					System.out.println("There is no current active session");
				else
				{			
					//request file from bob
					String fileReqName = Utils.GetUserInput("Enter filename to request: ");
					
					Message m = new Message(Message.MessageType.FileReceive);
					m.data = fileReqName.getBytes();
					
					SendMessage(m);
					m = ReceiveMessage();
					
					byte [] cipherBytes = m.data;
					EncryptClass decryptFile = new EncryptClass(cipher.toString(), encryptionKey, integrityKey, IV, cipherBytes, false);
					String fileSaveName = Utils.GetUserInput("Enter filename to save as: ");
					byte[] fileBytes = decryptFile.getText();
					Utils.writeFile(fileSaveName, fileBytes);
					
					in.close();
					out.close();
					socket.close();
				}
				
			}
			
			else if(opt.equals("5"))
			{
				socket = new Socket("localhost", 5555);
				out = new DataOutputStream(socket.getOutputStream());
				in = new DataInputStream(socket.getInputStream());
				
				Message m = new Message(Message.MessageType.SessionEstablished);
				SendMessage(m);
				break;//exit loop;
			}
			
			else
				System.out.println(opt + "is not a valid option");
		}
		
		in.close();
		out.close();
		socket.close();
	}
	
	public static boolean startSession() throws Exception
	{
		//values to remember
		byte[] e;
		byte[] n;
		byte[] bobR;
		byte[] aliceR;
		byte[] messagesToHash = new byte[32];
		Message m;
		
		//remember r sent to bob
		aliceR = IntToByte(getRandom());
		//send SessionStartRequest to server
		m = new Message(Message.MessageType.SessionStartRequest);
		m.r = aliceR;
		PrintToConsole(true, m);
		//save message for hashing later (store in first 8 bytes)
		System.arraycopy(m.ToBytes(), 0, messagesToHash, 0, 8);
		SendMessage(SimulateAttacker(m));

		//get responce from server
		m = ReceiveMessage();
		PrintToConsole(false, m);
		//set variables to use later
		sessionId = m.sessionId;
		bobR = m.r;
		e = m.e;
		n = m.n;
		cipher = m.cipher;
		//save message for hashing later (store in remaning bytes)
		System.arraycopy(m.ToBytes(), 0, messagesToHash, 8, 24);
		
		//remember s
		s = IntToByte(getRandom());
		//now that we have S, aliceR and bobR we can get K = f(S, aliceR, bobR)
		k = GenerateK(s, aliceR, bobR);
		
		//send PreMasterSecret
		m = new Message(Message.MessageType.PreMasterSecret);
		//encrypt s with bob public key before sending
		RsaEncryption rsa = new RsaEncryption(s, e, n);
		m.s = rsa.getCipherArray();
		//set keyed hash of messages
		m.hashOfMsgs = Utils.HMAC(k, messagesToHash);
		PrintToConsole(true, m);
		
		//create keyed hash that we should receive from bob
		byte[] h = new byte[56];
		System.arraycopy(messagesToHash, 0, h, 0, messagesToHash.length);
		System.arraycopy(m.ToBytes(), 0, h, messagesToHash.length, 24);	
		byte[] hash = Utils.HMAC(k, h);	
		
		SendMessage(SimulateAttacker(m));
		
		//get responce from server
		m = ReceiveMessage();
		PrintToConsole(false, m);
		
		if (m.type == Message.MessageType.KeyedHashOfMsgs)
		{
			//check received hash is correct and return result
			if(Arrays.equals(hash, m.hashOfMsgs))
			{
				m = new Message(Message.MessageType.SessionEstablished);
				PrintToConsole(true, m);
				SendMessage(m);
				return true;
			}
			else
			{
				//send Error to server
				m = new Message(Message.MessageType.Error);
				m.errorMessage = "Keyed hash of msgs it not correct. Session initiation failed";
				SendMessage(m);
				PrintToConsole(true, m);
			}
		}
		
		return false;
		
	}
	
	public static boolean resumeSession() throws Exception
	{
		//values to remember
		byte[] bobR;
		byte[] aliceR;
		byte[] messagesToHash = new byte[44];
		Message m;
		
		//remember r sent to bob
		aliceR = IntToByte(getRandom());
		//send SessionResumeRequest to server
		m = new Message(Message.MessageType.SessionResumeRequest);
		m.sessionId = sessionId;
		m.r = aliceR;
		PrintToConsole(true, m);
		//save message for hashing later (store in first 12 bytes)
		System.arraycopy(m.ToBytes(), 0, messagesToHash, 0, 12);
		SendMessage(SimulateAttacker(m));
		
		//get responce from server
		m = ReceiveMessage();
		PrintToConsole(false, m);
		if (m.type == Message.MessageType.SessionResumeResponce)
		{
			//set variables to use later
			sessionId = m.sessionId;
			bobR = m.r;
			cipher = m.cipher;
			//save message for hashing later (store in remaning bytes)
			System.arraycopy(m.ToBytes(), 0, messagesToHash, 12, 32);
			
			//now that we have S, aliceR and bobR we can get K = f(S, aliceR, bobR)
			k = GenerateK(s, aliceR, bobR);
			
			//create keyed hash that we should receive from bob
			byte[] hash = Utils.HMAC(k, Arrays.copyOfRange(messagesToHash, 0, 12));	
			
			//check received hash is correct
			if (Arrays.equals(hash, m.hashOfMsgs))
			{
				//send KeyedHashOfMsgs to server
				m = new Message(Message.MessageType.KeyedHashOfMsgs);
				//set keyed hash of messages
				m.hashOfMsgs = Utils.HMAC(k, messagesToHash);
				PrintToConsole(true, m);
				SendMessage(SimulateAttacker(m));
				
				try
				{
					m = ReceiveMessage();
				}
				catch(Exception e) { m = new Message(Message.MessageType.Error); }
				PrintToConsole(false, m);
				if (m.type == Message.MessageType.SessionEstablished)
					return true;
			}
			else
			{
				//send Error to server
				m = new Message(Message.MessageType.Error);
				m.errorMessage = "Keyed hash of msgs it not correct. Session initiation failed";
				SendMessage(m);
				PrintToConsole(true, m);
			}
		}
		
		return false;
	}
}